// Copyright (c) Microsoft Corporation. All rights reserved.

using System;
using System.IO;
using System.Runtime.InteropServices;
using System.Reflection;
using System.Runtime.Serialization;
using System.Threading;
using System.Text;
using Microsoft.Win32.SafeHandles;

namespace Microsoft.Tools.CodeRunner
{
	/// <summary>
	/// The exception that is thrown when one of the command line arguments provided 
	/// is not valid.
	/// </summary>
	[Serializable()]
	public sealed class RemoteEnvironmentException : ArgumentException
	{
		#region Public Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class.
		/// </summary>
		public RemoteEnvironmentException()
			: base()
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class
		/// with a descriptive message.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		public RemoteEnvironmentException(string message)
			: base(message)
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class
		/// with a descriptive message and an inner exception.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		/// <param name="innerException">A nested exception that is the cause of the current exception.</param>
		public RemoteEnvironmentException(string message, Exception innerException)
			: base(message, innerException)
		{
		}

		#endregion Public Instance Constructors

		#region Private Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="RemoteEnvironmentException" /> class 
		/// with serialized data.
		/// </summary>
		/// <param name="info">The <see cref="SerializationInfo" /> that holds the serialized object data about the exception being thrown.</param>
		/// <param name="context">The <see cref="StreamingContext" /> that contains contextual information about the source or destination.</param>
		private RemoteEnvironmentException(SerializationInfo info, StreamingContext context)
			: base(info, context)
		{
		}

		#endregion Private Instance Constructors
	}

	/// <summary>
	/// Summary description for RemoteEnvironment.
	/// </summary>
	public class RemoteEnvironment : IDisposable
	{
		#region Native methods

		[Flags]
		private enum PipeMode
		{
			Wait = 0x00000000,
			NoWait = 0x00000001,
			ReadModeByte = 0x00000000,
			ReadModeMessage = 0x00000002,
			TypeByte = 0x00000000,
			TypeMessage = 0x00000004,
		}

		private enum PipeAccess
		{
			Inbound = 0x00000001,
			Outbound = 0x00000002,
			Duplex = 0x00000003
		}
		
		private class Error
		{
			public const int PipeConnected = 535;
		}

		private class NativeMethods
		{
			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern bool SetNamedPipeHandleState(
				SafeFileHandle hPipe, ref PipeMode mode,
				IntPtr maxCollectionCount, IntPtr collectDataTimeout);

			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern bool WriteFile(
				SafeFileHandle hFile, byte[] buffer, int numberOfBytesToWrite,
				out int numberOfBytesWritten, IntPtr overlapped);

			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern SafeFileHandle CreateNamedPipe(
				string name, PipeAccess access, PipeMode mode, 
				int maxInstances, int outBufferSize, int inBufferSize,
				int defaultTimeout, IntPtr securityDescriptor);

			[DllImport("kernel32.dll", SetLastError = true)]
			public static extern bool ConnectNamedPipe(SafeFileHandle hPipe, IntPtr overlapped);
		}

		#endregion
	
		#region IDisposable and Related Members

		/// <summary>
		/// Finalizer
		/// </summary>
		~RemoteEnvironment()
		{
			Dispose(false);
		}
		
		/// <summary>
		/// Dispose method
		/// </summary>
		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}

		private void Dispose(bool disposing)
		{
			// If we are finalizing it's not safe to access any object members; this includes if an exception 
			// was thrown in the constructor.
			// If we are explicitly disposing then we assume that everything was initialized correctly if injector
			// is not null as that was the last thing created in the constructor.
			if (disposing && injector != null)
			{
				// First, we need to cause the DLL in the remote process to unload by sending a 
				// terminate message to the remote thread.
				byte[] bytes = Encoding.Unicode.GetBytes("terminate\0");
				int bytesWritten;
				
				NativeMethods.WriteFile(pipe, bytes, bytes.Length, out bytesWritten, IntPtr.Zero);
							
				// Now we are done with the pipe
				pipe.Close();
				pipe = null;

				// Wait forever for the remote thread to exit.  It's vital that it has exited before we release
				// the ProcessInjector object, or else the remote thread will return into freed memory and we 
				// bring down the remote process!  
				injector.WaitForExit();

				try
				{
					File.Delete(injector.DllFileName);
				}
				catch (SystemException)
				{
					// This is just a precaution, because nothing throws exceptions like I/O operations!
				}

				injector.Dispose();
				injector = null;
			}
		}
		#endregion	

		#region Public Methods
		/// <summary>
		/// Initializes a new instances of the <see cref="RemoteEnvironment"/> class.  Note that if this method 
		/// fails there may be a temporary file left in the TEMP directory.
		/// </summary>
		/// <param name="processId"></param>
		public RemoteEnvironment(int processId)
		{
			// Create a pipe to communicate with the remote thread.  This must be done first so that the 
			// pipe is ready when the remote thread starts
			pipe = NativeMethods.CreateNamedPipe(
				@"\\.\pipe\coderunner$", PipeAccess.Duplex, 
				PipeMode.TypeMessage | PipeMode.ReadModeMessage | PipeMode.Wait,
				255, maxMessageLength * Marshal.SizeOf(typeof(char)), maxMessageLength * Marshal.SizeOf(typeof(char)), 
				0, IntPtr.Zero);

			if (pipe.IsInvalid)
			{
				throw new RemoteEnvironmentException("Unable to create pipe to remote process");
			}

			string dllFileName = Path.GetTempFileName();

			// Unpack the DLL to be injected from our resources and put in the temporary file
			using (BinaryReader br = new BinaryReader(Assembly.GetExecutingAssembly().GetManifestResourceStream(
				typeof(RemoteEnvironment), "CodeRunnerHelper.dll")))
			{
				using (BinaryWriter bw = new BinaryWriter(new FileStream(dllFileName, FileMode.Create)))
				{
					bw.Write(br.ReadBytes((int)br.BaseStream.Length));
				}
			}

			// Inject the remote thread and pull in the DLL to create the other end of the pipe
			injector = new ProcessInjector(processId, dllFileName, "ControlRemoteEnvironment");

			// TODO: This will block until the injected thread connects. If anything goes wrong, such as the remote process dies, 
			// we'll hang. Let's not do that.
			bool connected = NativeMethods.ConnectNamedPipe(pipe, IntPtr.Zero) ?
				true : (Marshal.GetLastWin32Error() == Error.PipeConnected);
		}
		
		/// <summary>
		/// Set an environment variable in remote process
		/// </summary>
		/// <param name="name">Name of the environment variable</param>
		/// <param name="val">Value to set the environment variable to</param>
		public void SetEnvironmentVariable(string name, string val)
		{
			if (pipe == null)
				throw new InvalidOperationException("Object has been disposed");

			// Check maximum length of environment variable and name. The value includes space for the "set" command.
			if (name.Length + val.Length >= 32760)
				throw new ArgumentException("Length of name and value must be < 32760");
			
			// TODO: Security - AllowPartiallyTrusted Callers
		
			// Send the message to the remote thread
			string msg = String.Format("set\r\n{0}={1}\0", name, val);
			byte[] bytes = Encoding.Unicode.GetBytes(msg);
			int bytesWritten;

			NativeMethods.WriteFile(pipe, bytes, bytes.Length, out bytesWritten, IntPtr.Zero);
		}
		#endregion
		
		#region Private Data
		
		private ProcessInjector injector;
		private SafeFileHandle pipe;
		
		private const int maxMessageLength = 32768;
		
		#endregion
	}
}
